package club.janna.springboot.study.config;

import club.janna.springboot.study.util.AbstractArrayChooser;
import club.janna.springboot.study.util.Chooser;
import club.janna.springboot.study.util.JannaDataSourceHolder;
import lombok.extern.slf4j.Slf4j;
import org.springframework.jdbc.datasource.lookup.AbstractRoutingDataSource;

import javax.sql.DataSource;
import java.util.HashMap;
import java.util.Map;

@Slf4j
public class JannaReadWriteDataSource extends AbstractRoutingDataSource {

    private final String[] masterKeys;
    private final String[] slaveKeys;

    private final Chooser<String> masterChooser;
    private final Chooser<String> slaveChooser;

    public JannaReadWriteDataSource(DataSource[] masters, DataSource... slaves) {
        if(masters == null || masters.length == 0 || slaves == null || slaves.length == 0)
            throw new IllegalArgumentException("masters or slaves was null");

        masterKeys = new String[masters.length];
        slaveKeys = new String[slaves.length];

        Map<Object, Object> map = new HashMap<>();
        for (int i = 0; i < masters.length; i++) {
            masterKeys[i] = "Master-" + i;
            map.put(masterKeys[0], masters[0]);
        }
        for (int i = 0; i < slaves.length; i++) {
            slaveKeys[i] = "Slave-" + i;
            map.put(slaveKeys[i], slaves[i]);
        }
        setTargetDataSources(map);
        setDefaultTargetDataSource(masters[0]);

        masterChooser = masterKeys.length == 1 ? new AbstractArrayChooser.SingleArrayChooser<>(masterKeys) :
                (isPowerOfTwo(masterKeys.length) ?
                    new AbstractArrayChooser.PowerOfTwoArrayChooser<>(masterKeys) :
                    new AbstractArrayChooser.GenericArrayChooser<>(masterKeys));

        slaveChooser = slaveKeys.length == 1 ? new AbstractArrayChooser.SingleArrayChooser<>(slaveKeys) :
                (isPowerOfTwo(slaveKeys.length) ?
                    new AbstractArrayChooser.PowerOfTwoArrayChooser<>(slaveKeys) :
                    new AbstractArrayChooser.GenericArrayChooser<>(slaveKeys));
    }

    @Override
    protected Object determineCurrentLookupKey() {
        String key;
        if(JannaDataSourceHolder.get() == JannaDataSourceHolder.Type.READ)
            key = slaveChooser.next();
        else
            key = masterChooser.next();
        log.info("使用数据源 -> {}", key);
        return key;
    }

    private boolean isPowerOfTwo(int length) {
        return (length & -length) == length;
    }
}
